import numpy as np
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
from python_ai.common.xcommon import sep
from python_ai.ML.decomposition.my.x_3d_data import x_rotate_counter_clock
from python_ai.ML_2.decomposition.follow_teacher.x_pca_hyperbola_again import hyperbola


class PCAMyImpl(object):

    def __init__(self, n_components):
        self.n_components_ = n_components

    def fit_transform(self, X):
        m, n = X.shape

        X -= X.mean(axis=0)

        xcov = np.dot(X.T, X) / m  # ATTENTION "/ m" for Expection of (x-mu)**2 or (x1-mu1)(x2-mu2)

        lmds, us = np.linalg.eig(xcov)

        idx_sorted = lmds.argsort()[-1::-1]
        lmds = lmds[idx_sorted]  # (n, )
        us = us[idx_sorted]  # n x n
        lmds = lmds[:self.n_components_]  # (nc, )
        us = us[:self.n_components_]  # nc x n  # ATTENTION Result of np.linalg.eig is groups of unit vectors

        X_new = np.dot(X, us.T)
        self.components_ = us
        self.explained_variance_ = lmds
        return X_new


class PCAMyImplBySvd(object):

    def __init__(self, n_components):
        self.n_components_ = n_components

    def fit_transform(self, X):
        m, n = X.shape

        X -= X.mean(axis=0)

        xcov = np.cov(X.T)

        us, lmds, _ = np.linalg.svd(xcov)

        lmds = lmds[:self.n_components_]  # (nc, )
        us = us[:, :self.n_components_]  # n x nc  # ATTENTION Result of np.linalg.svd is matrix, [:, :nc] is the target unit vectors.

        X_new = np.dot(X, us)
        self.components_ = us.T
        self.explained_variance_ = lmds
        return X_new

if '__main__' == __name__:
    spr = 2
    spc = 4
    spn = 0
    plt.figure(figsize=[16, 8])
    cmap = plt.cm.get_cmap('rainbow', 2)

    # data
    n_dots = 16
    x_min = -6
    x_max = 6
    a = 2
    b = 3
    x_mu = 15
    y_mu = 30
    x = np.r_[np.linspace(x_min, -a, 1 + n_dots), np.linspace(a, x_max, 1 + n_dots)]
    x += x_mu
    x_min += x_mu
    x_max += x_mu

    # draw: hyperbola in xOy plane
    y = hyperbola(x, a, b, x_mu)
    z = np.zeros_like(x)
    y_min = -hyperbola(x_min, a, b, x_mu) + y_mu
    y_max = hyperbola(x_max, a, b, x_mu) + y_mu
    line_x = np.array([x_min, x_max, x_max, x_min, x_min])
    line_y = np.array([y_max, y_max, y_min, y_min, y_max])
    line_z = np.zeros_like(line_x)

    # draw: rotate 60 degree on axis y
    theta_by_y = np.pi / 3
    x2, z2 = x_rotate_counter_clock(x, z, theta_by_y)
    y2 = y
    line_x2, line_z2 = x_rotate_counter_clock(line_x, line_z, theta_by_y)
    line_y2 = line_y

    # draw: rotate 30 degree on axis z
    theta_by_z = np.pi / 6
    x3, y3 = x_rotate_counter_clock(x2, y2 + y_mu, theta_by_z)
    z3 = z2
    x3neg, y3neg = x_rotate_counter_clock(x2, -y2 + y_mu, theta_by_z)
    line_x3, line_y3 = x_rotate_counter_clock(line_x2, line_y2, theta_by_z)
    line_z3 = line_z2

    # data to decompose
    xyz_pos = np.c_[x3, y3, z3]
    xyz_neg = np.c_[x3neg, y3neg, z3]
    xyz = np.r_[xyz_pos, xyz_neg]


    def x_draw_hyperbola():
        global spn
        spn += 1
        ax = plt.subplot(spr, spc, spn, projection='3d')
        # draw: hyperbola in xOy plane
        ax.scatter3D(x, y + y_mu, z, s=10, c='g')
        ax.scatter3D(x, -y + y_mu, z, s=10, c='g')
        ax.plot3D(line_x, line_y, c='g')
        ax.set_xlabel('x')
        ax.set_ylabel('y')
        ax.set_zlabel('z')
        ax.grid()
        # draw: rotate 60 degree on axis y
        ax.scatter3D(x2, y2 + y_mu, z2, s=10, c='y')
        ax.scatter3D(x2, -y2 + y_mu, z2, s=10, c='y')
        ax.plot3D(line_x2, line_y2, line_z2, c='y')
        # draw: rotate 30 degree on axis z
        ax.scatter3D(x3, y3, z3, s=10, c='r')
        ax.scatter3D(x3neg, y3neg, z3, s=10, c='r')
        ax.plot3D(line_x3, line_y3, line_z3, c='r')
        return ax

    cmap = plt.cm.get_cmap('rainbow', 2)

    title = '3d=>2d by PCA'
    sep(title)
    ax = x_draw_hyperbola()
    ax.set_title(title)
    from sklearn.decomposition import PCA
    model = PCA(n_components=2)
    xy = model.fit_transform(xyz)
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.scatter(xy[:, 0], xy[:, 1])
    us = model.components_
    lmds = model.explained_variance_
    print(lmds)
    for i, u in enumerate(us):
        u *= lmds[i]
        ax.plot3D([0, u[0]], [0, u[1]], [0, u[2]], color=cmap(i))

    title = '3d=>2d by TruncatedSVD'
    sep(title)
    ax = x_draw_hyperbola()
    ax.set_title(title)
    from sklearn.decomposition import TruncatedSVD
    xyz_m = xyz - xyz.mean(axis=0)
    model = TruncatedSVD(n_components=2)
    xy = model.fit_transform(xyz_m)
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.scatter(xy[:, 0], xy[:, 1])
    us = model.components_
    lmds = model.explained_variance_
    print(lmds)
    for i, u in enumerate(us):
        u *= lmds[i]
        ax.plot3D([0, u[0]], [0, u[1]], [0, u[2]], color=cmap(i))


    title = '3d=>2d by self'
    sep(title)
    ax = x_draw_hyperbola()
    ax.set_title(title)
    model = PCAMyImpl(n_components=2)
    xy = model.fit_transform(xyz)
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.scatter(xy[:, 0], xy[:, 1])
    us = model.components_
    lmds = model.explained_variance_
    print(lmds)
    for i, u in enumerate(us):
        u *= lmds[i]
        ax.plot3D([0, u[0]], [0, u[1]], [0, u[2]], color=cmap(i))

    title = '3d=>2d by self svd'
    sep(title)
    ax = x_draw_hyperbola()
    ax.set_title(title)
    model = PCAMyImplBySvd(n_components=2)
    xy = model.fit_transform(xyz)
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.scatter(xy[:, 0], xy[:, 1])
    us = model.components_
    lmds = model.explained_variance_
    print(lmds)
    for i, u in enumerate(us):
        u *= lmds[i]
        ax.plot3D([0, u[0]], [0, u[1]], [0, u[2]], color=cmap(i))

    plt.show()
